# -*- coding: utf-8 -*-
"""
参数 Schema 解析器 - 将简化的参数描述转换为结构化的 schema

@author: AI Assistant
"""
import re
from typing import Dict, Any, Optional, Tuple
from loguru import logger


class SchemaParser:
    """Schema 解析器 - 解析参数描述字符串"""
    
    @staticmethod
    def parse_description(description: str) -> Dict[str, Any]:
        """解析参数描述字符串
        
        支持的格式：
        - "参数说明 - 范围: 0-5"
        - "参数说明 (单位) - 范围: 0-1"
        - "参数说明 (单位)"
        
        Args:
            description: 参数描述字符串
            
        Returns:
            解析后的参数定义字典
        """
        result = {
            "type": "float",
            "description": description,
            "unit": "",
            "min": None,
            "max": None
        }
        
        # 提取单位（括号内的内容）
        unit_match = re.search(r'\(([^)]+)\)', description)
        if unit_match:
            result["unit"] = unit_match.group(1)
        
        # 提取范围
        range_match = re.search(r'范围[：:]\s*(\d+\.?\d*)\s*[-~]\s*(\d+\.?\d*)', description)
        if range_match:
            result["min"] = float(range_match.group(1))
            result["max"] = float(range_match.group(2))
        
        # 提取描述（去除范围和单位信息）
        desc_text = description
        # 移除范围部分
        desc_text = re.sub(r'\s*-\s*范围[：:].*$', '', desc_text)
        # 移除单位部分
        desc_text = re.sub(r'\s*\([^)]+\)\s*', ' ', desc_text).strip()
        
        result["description"] = desc_text
        
        return result
    
    @staticmethod
    def convert_simple_schema_to_structured(simple_schema: Dict[str, Any]) -> Dict[str, Any]:
        """将简化的 schema 转换为结构化的 schema
        
        输入格式：
        {
            "parameters": {
                "A": "蓄水空间分布参数 - 范围: 0-5",
                "B": "下渗率形状系数 - 范围: 1-3"
            },
            "initialState": {
                "SA0": "初始表层蓄水量 (mm)"
            }
        }
        
        输出格式：
        {
            "parameters": {
                "A": {
                    "type": "float",
                    "min": 0,
                    "max": 5,
                    "unit": "",
                    "description": "蓄水空间分布参数"
                }
            },
            "initial_state": {
                "SA0": {
                    "type": "float",
                    "min": None,
                    "max": None,
                    "unit": "mm",
                    "description": "初始表层蓄水量"
                }
            }
        }
        
        Args:
            simple_schema: 简化格式的 schema
            
        Returns:
            结构化的 schema
        """
        structured_schema = {}
        
        # 处理 parameters
        if "parameters" in simple_schema:
            structured_schema["parameters"] = {}
            for param_name, param_desc in simple_schema["parameters"].items():
                structured_schema["parameters"][param_name] = SchemaParser.parse_description(param_desc)
        
        # 处理 initialState（转换为 initial_state）
        if "initialState" in simple_schema:
            structured_schema["initial_state"] = {}
            for param_name, param_desc in simple_schema["initialState"].items():
                structured_schema["initial_state"][param_name] = SchemaParser.parse_description(param_desc)
        
        logger.info(f"成功转换 schema - parameters: {len(structured_schema.get('parameters', {}))}, "
                   f"initial_state: {len(structured_schema.get('initial_state', {}))}")
        
        return structured_schema
    
    @staticmethod
    def load_and_convert_schema(schema_dict: Dict[str, Any]) -> Dict[str, Any]:
        """加载并转换 schema（便捷方法）
        
        Args:
            schema_dict: 简化格式的 schema 字典
            
        Returns:
            结构化的 schema
        """
        return SchemaParser.convert_simple_schema_to_structured(schema_dict)


# 使用示例
if __name__ == "__main__":
    import json
    from pathlib import Path
    
    # 加载示例 schema
    schema_file = Path(__file__).parent.parent / "docs" / "para_schema.json"
    
    if schema_file.exists():
        with open(schema_file, 'r', encoding='utf-8') as f:
            simple_schema = json.load(f)
        
        print("=" * 80)
        print("原始 Schema (简化格式):")
        print("=" * 80)
        print(json.dumps(simple_schema, ensure_ascii=False, indent=2))
        
        # 转换
        structured_schema = SchemaParser.convert_simple_schema_to_structured(simple_schema)
        
        print("\n" + "=" * 80)
        print("转换后的 Schema (结构化格式):")
        print("=" * 80)
        print(json.dumps(structured_schema, ensure_ascii=False, indent=2))
        
        print("\n" + "=" * 80)
        print("参数统计:")
        print("=" * 80)
        print(f"模型参数数量: {len(structured_schema.get('parameters', {}))}")
        print(f"初始状态参数数量: {len(structured_schema.get('initial_state', {}))}")
        
        # 显示几个示例参数
        print("\n参数示例:")
        for category in ["parameters", "initial_state"]:
            if category in structured_schema:
                print(f"\n{category}:")
                for i, (name, info) in enumerate(list(structured_schema[category].items())[:3]):
                    print(f"  - {name}:")
                    print(f"    描述: {info['description']}")
                    print(f"    范围: [{info['min']}, {info['max']}] {info['unit']}")
    else:
        print(f"❌ 找不到文件: {schema_file}")
        
        # 使用内置示例
        simple_schema = {
            "parameters": {
                "A": "蓄水空间分布参数 - 范围: 0-5",
                "K": "表层出流系数 - 范围: 0.1-0.9"
            },
            "initialState": {
                "SA0": "初始表层蓄水量 (mm)"
            }
        }
        
        print("使用内置示例:")
        structured_schema = SchemaParser.convert_simple_schema_to_structured(simple_schema)
        print(json.dumps(structured_schema, ensure_ascii=False, indent=2))

